from mujoco_py import MjViewer, const
import glfw

from ..planner import ManualPlanner


class Viewer(MjViewer):
    def __init__(self, sim, planner: ManualPlanner):
        super().__init__(sim)
        self.planner = planner

    def render(self):
        super().render()
        self.__add_extra_menu()

    def key_callback(self, window, key, scancode, action, mods):
        if key == glfw.KEY_W:
            self.planner.left_forward()
        elif key == glfw.KEY_S:
            self.planner.left_backward()
        elif key == glfw.KEY_A:
            self.planner.left_left()
        elif key == glfw.KEY_D:
            self.planner.left_right()
        elif key == glfw.KEY_Q:
            self.planner.left_up()
        elif key == glfw.KEY_E:
            self.planner.left_down()
        elif key == glfw.KEY_I:
            self.planner.right_forward()
        elif key == glfw.KEY_K:
            self.planner.right_backward()
        elif key == glfw.KEY_J:
            self.planner.right_left()
        elif key == glfw.KEY_L:
            self.planner.right_right()
        elif key == glfw.KEY_U:
            self.planner.right_up()
        elif key == glfw.KEY_O:
            self.planner.right_down()

    def __add_extra_menu(self):
        self.add_overlay(const.GRID_TOPRIGHT, "Forward      ", "[W]/[I]     ")
        self.add_overlay(const.GRID_TOPRIGHT, "Backward     ", "[S]/[K]     ")
        self.add_overlay(const.GRID_TOPRIGHT, "Left         ", "[A]/[J]     ")
        self.add_overlay(const.GRID_TOPRIGHT, "Right        ", "[D]/[L]     ")
        self.add_overlay(const.GRID_TOPRIGHT, "Up           ", "[Q]/[U]     ")
        self.add_overlay(const.GRID_TOPRIGHT, "Down         ", "[E]/[O]     ")

    def _create_full_overlay(self):
        pass
